library(LDlinkR)
library(LDheatmap)
library(tidyverse)

# store personal token in R environment
# !/bin/bash sudo vim ~/.Renviron
api_token <- Sys.getenv("LDLINK_TOKEN")

ld_candidate <- read_tsv('Archive/covid19/results/vindija_LD_candidate.txt') |>
  distinct(Existing_variation, .keep_all = TRUE) |>
  select(c(Location, BIOTYPE, Existing_variation)) |>
  mutate(position = str_extract(Location, '(?<=:).+(?=-)'),
         position = str_glue('chr14:{position}'))

# smaller LD block of ~60 kb containing 13 SNPs
block_candidate <- ld_candidate[-c(1:7),]

block_candidate |>
  mutate(numeric_pos = str_extract(Location, '(?<=:).+(?=-)') |>
           as.numeric()) |>
  write_csv('Archive/covid19/results/final_block_candidate13.csv')

# if also in Chayrskaya?
archaic4 <- read_delim('Archive/covid19/data/4archaic_IGHG1_LD.txt')

block_candidate |>
  mutate(numeric_pos = str_extract(Location, '(?<=:).+(?=-)') |>
           as.numeric()) |>
  left_join(archaic4, join_by(numeric_pos == position))


# randomly selected ctrl sites
sites_as_ctrl <- read_csv('Archive/covid19/results/CDX_LD_ctrl_sites.csv') |>
  mutate(position = str_glue('chr14:{POS}'))

# high quality sites --------
high_quality_fst <- read_csv('Archive/covid19/results/CDX_quality_Fst.csv')

high_qual10 <- slice_max(high_quality_fst, Fst, n = 10)

high_qual_ctrl <- high_quality_fst |>
  filter(position < min(high_qual10$position) | position > max(high_qual10$position)) |>
  slice_sample(n = 10)

# prune sites with 0 MAF or 0 D'
high_qual_rs <- high_quality_fst |>
  mutate(chrom_position = str_glue('chr14:{position}')) |>
  pull(chrom_position) |>
  append('chr14:106204113') |>
  SNPclip(token = api_token, pop = 'EAS') |>
  filter(!is.na(Position) & str_detect(Details, 'LD|kept')) |>
  mutate(numeric_pos = str_remove(Position, 'chr14:') |>
           as.numeric()) |>
  arrange(numeric_pos)

high_qual_ctrl <- high_qual_rs |>
  filter(!(numeric_pos %in% high_qual10$position)) |>
  slice_sample(n = 10)

rs_ctrl <- SNPclip(c(sites_as_ctrl$position, block_candidate$position),
                   token = api_token) |>
  filter(!is.na(Position)) |>
  mutate(numeric_pos = str_remove(Position, 'chr14:') |>
           as.numeric()) |>
  arrange(numeric_pos)

g2r_ldr2 <- LDmatrix(snps = rs_ctrl$RS_Number,
         pop = 'EAS',
         token = api_token) |>
  filter(RS_number == 'rs117518546') |>
  pivot_longer(cols = where(is.numeric),
               names_to = 'linked_snp',
               values_to = 'r2') |>
  select(-RS_number)

g2r_lddp <- LDmatrix(snps = rs_ctrl$RS_Number,
                     pop = 'EAS',
                     r2d = 'd',
                     token = api_token) |>
  filter(RS_number == 'rs117518546') |>
  pivot_longer(cols = where(is.numeric),
               names_to = 'linked_snp',
               values_to = 'd_prime') |>
  select(-RS_number)

# LD plot (heatmap) ---------
ld_matrix <- LDmatrix(snps = rs_ctrl$Position,
                     pop = 'EAS',
                     r2d = 'd',
                     token = api_token) |>
  column_to_rownames('RS_number') |>
  mutate(across(everything(), \(x)replace_na(x, 0L)))
  
pos_matrix <- ld_matrix |>
  colnames() |>
  SNPclip(pop = 'EAS', token = api_token) |>
  mutate(numeric_pos = str_remove(Position, 'chr14:') |>
           as.numeric())

ld_matrix |>
  as.matrix() |>
  LDheatmap(genetic.distances = pos_matrix$numeric_pos,
            flip = TRUE,
            LDmeasure = "D'",
            title = '20 SNP in IGHG1-G396R LD block',
            SNP.name = 'rs117518546',
            color = colorRampPalette(c('red','yellow','white'))(20)) |>
  LDheatmap.highlight(27,39, lwd = 3)

# high quality LD heatmap ----------
hq_ld_matrix <- high_qual_rs$Position |>
  LDmatrix(pop = 'EAS',
           r2d = 'd',
           token = api_token) |>
  column_to_rownames('RS_number') |>
  mutate(across(everything(), \(x)replace_na(x, 0L)))

hq_pos_matrix <- hq_ld_matrix |>
  colnames() |>
  SNPclip(pop = 'EAS', token = api_token) |>
  mutate(numeric_pos = str_remove(Position, 'chr14:') |>
           as.numeric())

block_index1 <- hq_pos_matrix$numeric_pos %in% high_qual10$position |>
  detect_index(isTRUE)

block_index2 <- hq_pos_matrix$numeric_pos %in% high_qual10$position |>
  detect_index(isTRUE, .dir = 'backward')

hq_ld_matrix |>
  as.matrix() |>
  LDheatmap(genetic.distances = hq_pos_matrix$numeric_pos,
            flip = TRUE,
            LDmeasure = "D'",
            title = '10 high-quality SNP in IGHG1-G396R LD block',
            SNP.name = 'rs117518546',
            color = colorRampPalette(c('red','yellow','white'))(20)) |>
  LDheatmap.highlight(block_index1, block_index2, lwd = 3)

g2r_linked <- g2r_lddp |>
  left_join(g2r_ldr2) |>
  left_join(ld_candidate, by = join_by(linked_snp == Existing_variation)) |>
  mutate(POS = as.numeric(position)) |>
  select(-position)

read_csv('Archive/covid19/results/CDX_most_divergent_SNP.csv',
         col_types = cols(position = col_character())) |>
  left_join(g2r_linked) |>
  arrange(POS) |>
  write_csv('Archive/covid19/results/CDX_most_linked_SNP.csv')


# find traits reported from IGHG1-GR and its linked SNP -------
# from GWAS catalog
associ_trait <- LDtrait('rs1050501',
        pop = "EAS",
        token = api_token)

write_csv(associ_trait, 'Archive/covid19/results/IGHG1_linked_snp_gwas_traits.csv')

# new sites after Fst computation
new_fst <- c(
  106188628,
  106257610,
  106156655,
  106243847,
  106204113
) |>
  map_chr(\(x)str_glue('chr14:{x}'))

new_links_dprime <- new_fst |>
  LDmatrix(pop = 'EAS',
         r2d = 'd',
         token = api_token) |>
  filter(RS_number == 'rs117518546') |>
  pivot_longer(cols = where(is.numeric),
               names_to = 'linked_snp',
               values_to = 'd_prime') |>
  select(-RS_number)

new_links_r2 <- new_fst |>
  LDmatrix(pop = 'EAS',
           r2d = 'r2',
           token = api_token) |>
  filter(RS_number == 'rs117518546') |>
  pivot_longer(cols = where(is.numeric),
               names_to = 'linked_snp',
               values_to = 'r2') |>
  select(-RS_number)

# use snpclip to get rs_number and position pair
new_fst |>
  LDlinkR::SNPclip(pop = 'EAS',
           token = api_token) |>
  arrange(Position)

# some IGHG4-ins structural variant
browning_sv <- read_csv('Archive/covid19/data/browning_sv.csv')

vitenan_ins <- browning_sv |>
  filter(str_detect(ID, 'HG02059_chr14') & SVTYPE == 'INS')
